from unicodedata import category
from PIL import Image
import numpy as np
from torch._C import _enter_dual_level

from .build import TransformLoader


@TransformLoader.register()
class LoadImage:
    def __init__(self, transform_config):
        pass

    def __call__(self, data):
        image = Image.open(data['path'])
        data['image'] = data
        return data


@TransformLoader.register()
class PILImageToNumpy:
    def __init__(self, transform_config):
        pass

    def __call__(self, data):
        image = np.asarray(data['image'])
        if len(image.shape) > 2:
            image = image[:, :, :3]
        data['image'] = image
        return data


@TransformLoader.register()
class ZScoreNormalize:
    def __init__(self, transform_config):
        self.pixel_mean = np.array(transform_config['pixel_mean'])
        self.pixel_std = np.array(transform_config['pixel_std'])

    def __call__(self, data):
        image = data['image']

        if len(image.shape) == 3:
            pixel_mean = self.pixel_mean.reshape(1, 1, -1)
            pixel_std = self.pixel_std.reshape(1, 1, -1)
        else:
            pixel_mean = self.pixel_mean.reshape(1, 1)
            pixel_std = self.pixel_std.reshape(1, 1)
        image = (image - pixel_mean) / pixel_std

        data['image'] = image
        return data


@TransformLoader.register()
class ExtractOneChannel:
    def __init__(self, transform_config):
        self.channel_id = transform_config['channel_id']

    def __call__(self, data):
        image = data['image']
        if len(image.shape) == 3:
            image = image[:, :, self.channel_id]
        data['image'] = image
        return data


@TransformLoader.register()
class RandomFlip:
    def __init__(self, transform_config):
        self.horizontal_probability = transform_config['horizontal_probability']
        self.vertical_probability = transform_config['vertical_probability']

    def __call__(self, data):
        image = data['image']
        h, w = image.shape[:2]

        if np.random.rand() < self.horizontal_probability:
            image = image[:, ::-1]
            for obj in data['objects']:
                obj['box'][0], obj['box'][2] = w - obj['box'][2], w - obj['box'][0]

        if np.random.rand() < self.vertical_probability:
            image = image[::-1, :]
            for obj in data['objects']:
                obj['box'][1], obj['box'][3] = h - obj['box'][3], h - obj['box'][1]

        data['image'] = image
        return data


@TransformLoader.register()
class RandomApply:
    def __init__(self, transform_config):
        self.transform = TransformLoader.get(transform_config['transform'])
        self.probability = transform_config['probability']

    def __call__(self, data):
        if np.random.rand() < self.probability:
            data = self.transform(data)
        return data


@TransformLoader.register()
class RandomCrop:
    def __init__(self, transform_config):
        self.min_height_ratio = transform_config['min_height_ratio']
        self.min_width_ratio = transform_config['min_width_ratio']

    def __call__(self, data):
        image = data['image']

        h, w = image.shape[:2]

        height_ratio = np.random.rand() * (1.0 - self.min_height_ratio) + self.min_height_ratio
        width_ratio = np.random.rand() * (1.0 - self.min_width_ratio) + self.min_width_ratio

        new_h = int(h * height_ratio)
        new_w = int(w * width_ratio)

        x = int(np.random.rand() * (w - new_w))
        y = int(np.random.rand() * (h - new_h))

        image = image[y : y + new_h, x : x + new_w]

        clamp = lambda x, lower, upper: min(max(x, lower), upper)

        new_objects = []
        for obj in data['objects']:
            inter_w = max(min(x + new_w, obj['box'][2]) - max(x, obj['box'][0], 0))
            inter_h = max(min(y + new_h, obj['box'][3]) - max(y, obj['box'][1], 0))
            inter_area_ratio = inter_h * inter_w / ((obj['box'][3] - obj['box'][1]) * (obj['box'][2] - obj['box'][0]))
            if inter_area_ratio < 0.5:
                continue

            obj['box'][0] = clamp(obj['box'][0] - x, 0., new_w)
            obj['box'][1] = clamp(obj['box'][1] - y, 0., new_h)
            obj['box'][2] = clamp(obj['box'][2] - x, 0., new_w)
            obj['box'][3] = clamp(obj['box'][3] - y, 0., new_h)
            new_objects.append(obj)

        data['image'] = image
        data['objects'] = new_objects

        return data


@TransformLoader.register()
class MapCategoryToId:
    def __init__(self, transform_config):
        self.category_id_map = {category : i for i, category in enumerate(transform_config['categories'])}

    def __call__(self, data):
        for obj in data['objects']:
            obj['category_id'] = self.category_id_map[obj['category']]
        return data


@TransformLoader.register()
class ResizePadding:
    def __init__(self, transform_config):
        self.size = transform_config['size']

    def __call__(self, data):
        image = data['image']

        width, height = self.size
        old_height, old_width = image.shape[:2]
        if old_height > old_width:
            scale_ratio = width / old_width
            new_width = width
            new_height = int(old_height * scale_ratio)
        else:
            scale_ratio = height / old_height
            new_width = int(old_width * scale_ratio)
            new_height = height

        canvas = np.zeros((height, width) + image.shape[2:], np.float32)

        if (new_width, new_height) != (old_height, old_width):
            canvas[:new_height, :new_width] = cv2.resize(image, (new_width, new_height))
        else:
            canvas = image

        for obj in data['objects']:
            obj['box'] = [value * scale_ratio for value in obj['box']]

        data['image'] = canvas
        data['params']['resize_padding'] = {"scale_ratio": scale_ratio}
